import matplotlib.pyplot as plt
import matplotlib
import pandas as pd
import os
from datetime import datetime

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
matplotlib.use('Agg')  # 使用非交互式后端


class SimpleTradingVisualizer:
    # 简单的交易可视化类

    def __init__(self, equity_curve, trade_history, signals, price_data):
        self.equity_curve = equity_curve
        self.trade_history = trade_history
        self.signals = signals
        self.price_data = price_data
        # 创建保存图片的目录
        self.output_dir = self.create_output_directory()

    def create_output_directory(self):
        # 创建输出目录
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = f"trading_results_{timestamp}"
        os.makedirs(output_dir, exist_ok=True)
        return output_dir

    def save_plot(self, filename):
        # 保存当前图表到文件
        filepath = os.path.join(self.output_dir, filename)
        plt.savefig(filepath, dpi=300, bbox_inches='tight', facecolor='white')
        print(f"图表已保存: {filepath}")
        plt.close()  # 关闭图表以释放内存

    def plot_equity_curve(self):
        # 简单的累积收益曲线
        if not self.equity_curve:
            print("没有权益数据")
            return

        equity_df = pd.DataFrame(self.equity_curve)

        plt.figure(figsize=(12, 6))
        plt.plot(equity_df['date'], equity_df['equity'], 'b-', linewidth=2, label='账户资金')
        plt.title('账户资金变化曲线')
        plt.xlabel('日期')
        plt.ylabel('资金 (元)')
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.xticks(rotation=45)
        plt.tight_layout()
        self.save_plot('equity_curve.png')

    def plot_win_rate(self):
        # 简单的胜率图表
        if not self.trade_history:
            print("没有交易记录")
            return

        trades_df = pd.DataFrame(self.trade_history)

        # 计算胜率
        winning_trades = len(trades_df[trades_df['pnl'] > 0])
        losing_trades = len(trades_df[trades_df['pnl'] <= 0])
        total_trades = len(trades_df)
        win_rate = (winning_trades / total_trades) * 100 if total_trades > 0 else 0

        # 创建图表
        plt.figure(figsize=(12, 6))

        # 盈亏分布饼图
        plt.subplot(1, 2, 1)
        labels = ['盈利交易', '亏损交易']
        sizes = [winning_trades, losing_trades]
        colors = ['lightgreen', 'lightcoral']
        plt.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', startangle=90)
        plt.title(f'交易胜率: {win_rate:.1f}%')

        # 多空交易柱状图
        plt.subplot(1, 2, 2)
        if 'trade_type' in trades_df.columns:
            long_trades = len(trades_df[trades_df['trade_type'] == 'long'])
            short_trades = len(trades_df[trades_df['trade_type'] == 'short'])

            plt.bar(['做多', '做空'], [long_trades, short_trades],
                    color=['lightblue', 'lightpink'], alpha=0.7)
            plt.title('多空交易次数')
            plt.ylabel('交易次数')

        plt.tight_layout()
        self.save_plot('win_rate_analysis.png')

        # 创建详细的交易统计图表
        self.plot_detailed_statistics(trades_df)

    def plot_detailed_statistics(self, trades_df):
        # 盈亏分布直方图
        plt.figure(figsize=(12, 10))

        # 1. 盈亏分布直方图
        plt.subplot(2, 2, 1)
        plt.hist(trades_df['pnl'], bins=20, color='skyblue', edgecolor='black', alpha=0.7)
        plt.axvline(x=0, color='red', linestyle='--', linewidth=1)
        plt.title('盈亏分布直方图')
        plt.xlabel('盈亏金额')
        plt.ylabel('交易次数')
        plt.grid(True, alpha=0.3)

        # 2. 收益率分布
        plt.subplot(2, 2, 2)
        if 'return_pct' in trades_df.columns:
            plt.hist(trades_df['return_pct'], bins=20, color='lightcoral', edgecolor='black', alpha=0.7)
            plt.axvline(x=0, color='red', linestyle='--', linewidth=1)
            plt.title('收益率分布 (%)')
            plt.xlabel('收益率 (%)')
            plt.ylabel('交易次数')
            plt.grid(True, alpha=0.3)

        # 4. 多空交易盈亏对比
        plt.subplot(2, 2, 4)
        if 'trade_type' in trades_df.columns:
            long_pnl = trades_df[trades_df['trade_type'] == 'long']['pnl'].mean()
            short_pnl = trades_df[trades_df['trade_type'] == 'short']['pnl'].mean()

            plt.bar(['做多', '做空'], [long_pnl, short_pnl],
                    color=['lightblue', 'lightpink'], alpha=0.7)
            plt.title('多空交易平均盈亏对比')
            plt.ylabel('平均盈亏金额')
            plt.grid(True, alpha=0.3)

        plt.tight_layout()
        self.save_plot('detailed_statistics.png')

    def plot_trading_signals(self):
        # 简单的多空点位标注图
        if not self.signals or self.price_data is None:
            print("没有信号数据或价格数据")
            return

        signals_df = pd.DataFrame(self.signals)
        price_df = self.price_data.copy()

        plt.figure(figsize=(14, 8))

        # 绘制价格曲线
        plt.plot(price_df.index, price_df['close'], 'k-', linewidth=1, label='收盘价', alpha=0.7)

        # 分离开仓和平仓信号
        entry_signals = signals_df[signals_df['type'] == 'entry']
        exit_signals = signals_df[signals_df['type'] == 'exit']

        # 绘制做多信号
        long_entries = entry_signals[entry_signals['signal'] == 2]
        long_exits = exit_signals[exit_signals['signal'] == 2]

        if not long_entries.empty:
            plt.scatter(long_entries['date'], long_entries['price'],
                        color='green', marker='^', s=80, label='做多开仓')
        if not long_exits.empty:
            plt.scatter(long_exits['date'], long_exits['price'],
                        color='darkgreen', marker='v', s=80, label='做多平仓')

        # 绘制做空信号
        short_entries = entry_signals[entry_signals['signal'] == 0]
        short_exits = exit_signals[exit_signals['signal'] == 0]

        if not short_entries.empty:
            plt.scatter(short_entries['date'], short_entries['price'],
                        color='red', marker='v', s=80, label='做空开仓')
        if not short_exits.empty:
            plt.scatter(short_exits['date'], short_exits['price'],
                        color='darkred', marker='^', s=80, label='做空平仓')

        plt.title('价格走势与交易信号')
        plt.xlabel('日期')
        plt.ylabel('价格')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.xticks(rotation=45)
        plt.tight_layout()
        self.save_plot('trading_signals.png')

    def plot_all(self):
        # 绘制所有图表并保存到文件
        print("生成交易可视化图表...")
        print(f"图表将保存到目录: {os.path.abspath(self.output_dir)}")

        self.plot_equity_curve()
        self.plot_win_rate()
        self.plot_trading_signals()

        print("所有图表已生成！")
